package 中等.搜索.广度优先搜索;

import java.util.*;

/**
 * 在一个有向图中，节点分别标记为 0, 1, ..., n-1。图中每条边为红色或者蓝色，
 * 且存在自环或平行边。
 * red_edges 中的每一个 [i, j] 对表示从节点 i 到节点 j 的红色有向边。类似
 * 地，blue_edges 中的每一个 [i, j] 对表示从节点 i 到节点 j 的蓝色有向边。
 * 返回长度为 n 的数组 answer，其中 answer[X] 是从节点 0 到节点 X 的红色
 * 边和蓝色边交替出现的最短路径的长度。如果不存在这样的路径，那么 answer[x] = -1。
 * <p>
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode.cn/problems/shortest-path-with-alternating-colors
 */
public class 颜色交替的最短路径_1129 {

    public static void main(String[] args) {

        颜色交替的最短路径_1129 impl = new 颜色交替的最短路径_1129();

        System.out.println(Arrays.toString(impl.shortestAlternatingPaths(5,
                new int[][]{{2, 2}, {0, 4}, {4, 2}, {4, 3}, {2, 4}, {0, 0}, {0, 1}, {2, 3}, {1, 3}},
                new int[][]{{0, 4}, {1, 0}, {1, 4}, {0, 0}, {4, 0}})));

    }

    /**
     * 广度优先搜索+模拟
     * 1，从 0 开始搜索，存在两种方案 红-蓝-红 或者 蓝-红-蓝
     * 所以需要搜索两次，取较小距离
     * 2，因为存在环，需要记录访问过的元素
     * 注意：这里需要分开记录，一个元素可以在红色中选择一次，和
     * 在蓝色中选择一次
     */
    public int[] shortestAlternatingPaths(int n, int[][] redEdges, int[][] blueEdges) {
        int[] ans = new int[n];
        Arrays.fill(ans, -1);
        Map<Integer, List<Integer>> redMap = array2Map(redEdges);
        Map<Integer, List<Integer>> blueMap = array2Map(blueEdges);

        bfs(ans, redMap, blueMap,
                true);
        bfs(ans, redMap, blueMap,
                false);
        return ans;
    }

    private void bfs(int[] ans, Map<Integer, List<Integer>> redMap,
                     Map<Integer, List<Integer>> blueMap, boolean initColor) {
        Set<Integer> redVisited = new HashSet<>(), blueVisited = new HashSet<>();
        redVisited.add(0);
        int distance = 0;
        Deque<Integer> queue = new ArrayDeque<>();
        queue.addLast(0);
        boolean needRed = initColor;
        while (!queue.isEmpty()) {
            int size = queue.size();
            while (size-- > 0) {
                Integer cur = queue.pollFirst();
                if (ans[cur] == -1) {
                    ans[cur] = distance;
                } else {
                    ans[cur] = Math.min(ans[cur], distance);
                }

                List<Integer> nextList;
                Set<Integer> visited;
                if (needRed) {
                    nextList = redMap.get(cur);
                    visited = redVisited;
                } else {
                    nextList = blueMap.get(cur);
                    visited = blueVisited;
                }
                if (nextList != null) {
                    for (Integer next : nextList) {
                        if (!visited.contains(next)) {
                            queue.addLast(next);
                            visited.add(next);
                        }
                    }
                }
            }
            needRed ^= true;
            distance++;
        }
    }

    private Map<Integer, List<Integer>> array2Map(int[][] edges) {
        Map<Integer, List<Integer>> map = new HashMap<>();
        for (int[] edge : edges) {
            List<Integer> list = map.get(edge[0]);
            if (list == null) {
                list = new ArrayList<>();
                map.put(edge[0], list);
            }
            list.add(edge[1]);
        }
        return map;
    }

}
